# 染色体非整倍体跨性别预测优化代码
import warnings
warnings.filterwarnings('ignore')
import util_for_output_zh

import os
import pandas as pd
import numpy as np
from util_set_zh_matplot import plt
import seaborn as sns
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.metrics import classification_report, confusion_matrix, f1_score
from sklearn.ensemble import RandomForestClassifier
import xgboost as xgb
from imblearn.over_sampling import SMOTE
from imblearn.pipeline import Pipeline
import joblib

# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)

# 创建输出目录
out_dir = 'ques4_prepare_data'
os.makedirs(out_dir, exist_ok=True)


def main():
    """主函数：控制流程"""
    # 执行主分析流程
    stage1()



def plot_category_distributions(female_data, male_data):
    """绘制男女及总体类别分布"""
    female_counts = female_data['染色体的非整倍体'].value_counts()
    male_counts = male_data['染色体的非整倍体'].value_counts()
    total_counts = female_counts.add(male_counts, fill_value=0).sort_values(ascending=False)

    plt.figure(figsize=(15, 12))
    
    plt.subplot(3, 1, 1)
    female_counts.plot(kind='bar', color='pink')
    plt.title('女胎染色体非整倍体类别分布')
    plt.ylabel('样本数')
    
    plt.subplot(3, 1, 2)
    male_counts.plot(kind='bar', color='lightblue')
    plt.title('男胎染色体非整倍体类别分布')
    plt.ylabel('样本数')
    
    plt.subplot(3, 1, 3)
    total_counts.plot(kind='bar', color='purple')
    plt.title('总体染色体非整倍体类别分布')
    plt.ylabel('样本数')
    
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, '类别分布统计.png'))
    plt.close()
    print("类别分布图表已保存")


def smart_data_augmentation(df, features, n_samples):
    """智能数据增强：基于特征分布生成合理的新样本"""
    augmented = []
    # 对每个特征，根据其分布生成新样本
    for i in range(n_samples):
        new_sample = {}
        for feature in features:
            # 获取特征分布参数
            mean = df[feature].mean()
            std = df[feature].std()
            # 基于正态分布生成新值（确保在合理范围内）
            min_val = df[feature].min()
            max_val = df[feature].max()
            
            # 生成在3倍标准差范围内的值
            while True:
                val = np.random.normal(mean, std)
                if min_val <= val <= max_val:
                    break
            new_sample[feature] = val
        augmented.append(new_sample)
    
    return pd.DataFrame(augmented)


def balance_dataset(data, features, label_col, max_samples=700):
    """平衡数据集：结合下采样和智能数据增强"""
    balanced = pd.DataFrame()
    labels = data[label_col].unique()
    
    for label in labels:
        label_data = data[data[label_col] == label].copy()
        current_count = len(label_data)
        
        if current_count >= max_samples:
            # 下采样
            balanced_subset = label_data.sample(n=max_samples, random_state=42)
        else:
            # 需要生成的样本数
            need = max_samples - current_count
            # 智能数据增强
            augmented = smart_data_augmentation(label_data, features, need)
            # 保留标签和性别信息
            augmented[label_col] = label
            augmented['性别'] = label_data['性别'].iloc[0]
            # 合并原始和增强数据
            balanced_subset = pd.concat([label_data, augmented], ignore_index=True)
        
        balanced = pd.concat([balanced, balanced_subset], ignore_index=True)
    
    return balanced


def cross_gender_test(data, features, le, train_gender, test_gender):
    """跨性别预测测试"""
    # 划分训练集和测试集
    train_data = data[data['性别'] == train_gender].copy()
    test_data = data[data['性别'] == test_gender].copy()
    
    # 提取特征和标签
    X_train = train_data[features]
    y_train = le.transform(train_data['染色体的非整倍体'])
    X_test = test_data[features]
    y_test = le.transform(test_data['染色体的非整倍体'])
    
    # 数据标准化
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    
    # 处理训练集中的类别不平衡
    if len(np.unique(y_train)) > 1:  # 确保有多个类别
        smote = SMOTE(random_state=42, k_neighbors=min(5, len(X_train)-1))
        X_train_resampled, y_train_resampled = smote.fit_resample(X_train_scaled, y_train)
    else:
        X_train_resampled, y_train_resampled = X_train_scaled, y_train
    
    # 构建更鲁棒的模型
    clf = xgb.XGBClassifier(
        objective='multi:softmax',
        num_class=len(le.classes_),
        n_estimators=200,
        max_depth=4,  # 减小深度防止过拟合
        learning_rate=0.05,
        subsample=0.8,  # 随机采样训练样本
        colsample_bytree=0.8,  # 随机采样特征
        reg_alpha=0.1,  # L1正则化
        reg_lambda=0.1,  # L2正则化
        scale_pos_weight=1,
        random_state=42
    )
    
    # 训练模型
    clf.fit(X_train_resampled, y_train_resampled)
    
    # 预测与评估
    y_pred = clf.predict(X_test_scaled)
    print(f"\n{train_gender}训练集预测{test_gender}测试集结果:")
    print(classification_report(y_test, y_pred, target_names=le.classes_))
    
    # 绘制混淆矩阵
    plot_confusion_matrix(y_test, y_pred, le.classes_, 
                         f"{train_gender}→{test_gender}混淆矩阵")
    
    # 交叉验证评估模型稳定性
    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    cv_scores = cross_val_score(clf, X_train_resampled, y_train_resampled, 
                               cv=cv, scoring='f1_macro')
    print(f"训练集交叉验证F1分数: {cv_scores.mean():.4f} ± {cv_scores.std():.4f}")
    
    return clf, scaler


def plot_confusion_matrix(y_true, y_pred, labels, title):
    """绘制混淆矩阵"""
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=labels, yticklabels=labels)
    plt.xlabel('预测类别')
    plt.ylabel('实际类别')
    plt.title(title)
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, f'{title}.png'))
    plt.close()


def analyze_feature_correlation(data, features, gender_col, label_col):
    """分析特征与性别、标签的相关性"""
    # 计算特征与性别的相关性（将性别转换为数值）
    data_corr = data.copy()
    data_corr[gender_col] = data_corr[gender_col].map({'男': 0, '女': 1})
    
    # 计算相关性
    corr_with_gender = data_corr[features + [gender_col]].corr()[gender_col].drop(gender_col)
    
    # 绘制相关性热图
    plt.figure(figsize=(12, 8))
    sns.barplot(x=corr_with_gender.index, y=corr_with_gender.values)
    plt.title('特征与性别的相关性')
    plt.xticks(rotation=45, ha='right')
    plt.axhline(y=0.3, color='r', linestyle='--', alpha=0.5, label='高相关性阈值')
    plt.legend()
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, '特征与性别的相关性.png'))
    plt.close()
    
    print("\n特征与性别的相关性（绝对值）:")
    print(corr_with_gender.abs().sort_values(ascending=False))
    
    return corr_with_gender


def stage1():
    """主分析流程"""
    # 加载数据
    file_path = 'ques4_prepare_data/ques4_concat.csv'
    combined_data = pd.read_csv(file_path)

    # 分离男女数据
    female_data = combined_data[combined_data['性别'] == '女'].copy()
    male_data = combined_data[combined_data['性别'] == '男'].copy()

    # 过滤有效类别（样本数≥10且男女均存在）
    total_counts = combined_data['染色体的非整倍体'].value_counts()
    valid_categories = [label for label in total_counts.index 
                       if total_counts[label] >= 10 
                       and label in female_data['染色体的非整倍体'].unique()
                       and label in male_data['染色体的非整倍体'].unique()]
    
    print(f"原始类别分布:\n{total_counts}\n")
    print(f"过滤后保留的类别:\n{valid_categories}\n")

    # 筛选数据
    final_data = combined_data[combined_data['染色体的非整倍体'].isin(valid_categories)]
    print(f"最终使用的共同类别:\n{valid_categories}\n")

    # 定义初始特征集（包含所有可用特征）
    initial_features = [
        '原始读段数', '在参考基因组上比对的比例', '重复读段的比例', '唯一比对的读段数',
        'GC含量', '13号染色体的Z值', '18号染色体的Z值', '21号染色体的Z值', 
        'X染色体的Z值', 'X染色体浓度', '13号染色体的GC含量', 
        '18号染色体的GC含量', '21号染色体的GC含量'
    ]

    # 分析特征与性别的相关性，识别性别敏感特征
    corr_with_gender = analyze_feature_correlation(
        final_data, initial_features, '性别', '染色体的非整倍体'
    )
    
    # 筛选低性别相关性特征（相关性绝对值 < 0.3）
    gender_agnostic_features = [
        f for f in initial_features 
        if abs(corr_with_gender.get(f, 0)) < 0.3
    ]
    
    # 如果筛选后特征太少，保留核心生物学特征
    if len(gender_agnostic_features) < 5:
        print("\n警告：低性别相关性特征较少，保留核心染色体特征")
        gender_agnostic_features = [
            '13号染色体的Z值', '18号染色体的Z值', '21号染色体的Z值',
            '13号染色体的GC含量', '18号染色体的GC含量', '21号染色体的GC含量'
        ]
    
    print(f"\n使用的低性别相关性特征:\n{gender_agnostic_features}")

    # 检查特征存在性
    missing_features = [f for f in gender_agnostic_features 
                       if f not in final_data.columns]
    if missing_features:
        raise ValueError(f"特征不存在: {missing_features}")

    # 按性别分别平衡数据
    balanced_data = pd.DataFrame()
    max_samples_per_class = 700  # 每个类别最大样本数
    
    for gender in ['女', '男']:
        gender_data = final_data[final_data['性别'] == gender].copy()
        balanced_gender = balance_dataset(
            gender_data, 
            gender_agnostic_features, 
            '染色体的非整倍体',
            max_samples_per_class
        )
        balanced_data = pd.concat([balanced_data, balanced_gender], ignore_index=True)

    # 打印平衡前后分布
    print("\n平衡前类别分布:")
    print(final_data['染色体的非整倍体'].value_counts())
    print("\n平衡后类别分布:")
    print(balanced_data['染色体的非整倍体'].value_counts())

    # 标签编码
    le = LabelEncoder()
    le.fit(valid_categories)  # 确保编码一致性

    # 跨性别测试
    clf_male2female, scaler_male = cross_gender_test(
        balanced_data, gender_agnostic_features, le, '男', '女'
    )
    clf_female2male, scaler_female = cross_gender_test(
        balanced_data, gender_agnostic_features, le, '女', '男'
    )

    # 特征重要性分析
    print("\n特征重要性排序:")
    importance = pd.DataFrame({
        'feature': gender_agnostic_features,
        'importance': clf_female2male.feature_importances_
    }).sort_values('importance', ascending=False)
    print(importance)

    # 保存模型和编码器
    joblib.dump({
        'model': clf_male2female,
        'scaler': scaler_male,
        'label_encoder': le,
        'features': gender_agnostic_features
    }, os.path.join(out_dir, 'male2female_model.pkl'))
    
    joblib.dump({
        'model': clf_female2male,
        'scaler': scaler_female,
        'label_encoder': le,
        'features': gender_agnostic_features
    }, os.path.join(out_dir, 'female2male_model.pkl'))
    
    print("\n模型已保存至 ques4_prepare_data 目录")


if __name__ == '__main__':
    main()
